#!/usr/bin/env python3
"""项目初始化脚本"""

import os
import sys
import subprocess
import platform
import shutil
from pathlib import Path

def check_python_version():
    """检查Python版本是否满足要求"""
    required_version = (3, 11)
    current_version = sys.version_info[:2]
    
    print(f"当前Python版本: {platform.python_version()}")
    
    if current_version < required_version:
        print(f"错误: 需要Python {required_version[0]}.{required_version[1]} 或更高版本")
        sys.exit(1)
    
    print(f"Python版本检查通过 ✅")
    return True



def create_virtual_env():
    """创建Python虚拟环境"""
    venv_dir = Path("venv")
    
    if venv_dir.exists():
        print(f"虚拟环境已存在于 {venv_dir}")
        return venv_dir
    
    print(f"正在创建虚拟环境...")
    
    try:
        subprocess.run(
            [sys.executable, "-m", "venv", venv_dir],
            check=True,
            capture_output=True,
            text=True
        )
        
        print(f"虚拟环境创建成功: {venv_dir} ✅")
        return venv_dir
    except subprocess.CalledProcessError as e:
        print(f"创建虚拟环境失败: {e.stderr}")
        sys.exit(1)


def install_dependencies(venv_dir):
    """安装项目依赖"""
    requirements_file = Path("requirements.txt")
    
    if not requirements_file.exists():
        print(f"错误: 找不到requirements.txt文件")
        sys.exit(1)
    
    # 确定pip路径
    if platform.system() == "Windows":
        pip_path = venv_dir / "Scripts" / "pip.exe"
    else:
        pip_path = venv_dir / "bin" / "pip"
    
    print(f"正在安装项目依赖...")
    
    try:
        # 更新pip
        subprocess.run(
            [str(pip_path), "install", "--upgrade", "pip"],
            check=True,
            capture_output=True,
            text=True
        )
        
        # 安装依赖
        subprocess.run(
            [str(pip_path), "install", "-r", str(requirements_file)],
            check=True,
            capture_output=True,
            text=True
        )
        
        print(f"依赖安装成功 ✅")
        return True
    except subprocess.CalledProcessError as e:
        print(f"安装依赖失败: {e.stderr}")
        print("\n提示: 如果PyTorch安装失败，可能需要根据你的CUDA版本手动安装合适的PyTorch版本")
        print("请参考: https://pytorch.org/get-started/locally/")
        sys.exit(1)


def create_required_directories():
    """创建项目所需的目录"""
    directories = [
        "config",
        "models",
        "fine_tune",
        "rag",
        "api",
        "examples",
        "rag_data",
        "rag_index",
        "finetune_output"
    ]
    
    print(f"正在创建必要的目录...")
    
    for dir_name in directories:
        dir_path = Path(dir_name)
        if not dir_path.exists():
            try:
                dir_path.mkdir(parents=True, exist_ok=True)
                print(f"创建目录: {dir_path}")
            except Exception as e:
                print(f"创建目录 {dir_path} 失败: {e}")
                sys.exit(1)
    
    print(f"目录创建完成 ✅")
    return True


def create_env_file():
    """创建.env文件（如果不存在）"""
    env_file = Path(".env")
    env_example_file = Path(".env.example")
    
    if env_file.exists():
        print(f".env文件已存在")
        return True
    
    if env_example_file.exists():
        print(f"正在从.env.example创建.env文件...")
        try:
            shutil.copy2(env_example_file, env_file)
            print(f".env文件创建成功 ✅")
            return True
        except Exception as e:
            print(f"创建.env文件失败: {e}")
            sys.exit(1)
    else:
        print(f"警告: 找不到.env.example文件，跳过创建.env文件")
        return False


def main():
    """主函数"""
    print("=== AI大模型二次开发项目初始化脚本 ===")
    
    # 1. 检查Python版本
    check_python_version()
    
    # 2. 创建虚拟环境
    venv_dir = create_virtual_env()
    
    # 3. 安装依赖
    install_dependencies(venv_dir)
    
    # 4. 创建必要的目录
    create_required_directories()
    
    # 5. 创建.env文件
    create_env_file()
    
    print("\n=== 项目初始化完成 ===")
    print(f"\n下一步操作:")
    print(f"1. 激活虚拟环境:")
    if platform.system() == "Windows":
        print(f"   venv\\Scripts\\activate")
    else:
        print(f"   source venv/bin/activate")
    print(f"2. 启动服务:")
    print(f"   python main.py")
    print(f"3. 访问API文档:")
    print(f"   http://localhost:8000/docs")
    print(f"4. 查看前端示例:")
    print(f"   打开 examples/web_interface.html 文件")
    print(f"\n注意: 首次运行时，系统会自动下载Qwen2.5-7B-Instruct模型，这可能需要一些时间和磁盘空间")


if __name__ == "__main__":
    main()